use std::{collections::HashMap, path::Path, usize};

use melior::{
    Context,
    ir::{
        Attribute, Block, BlockLike, BlockRef, Identifier, Location, Module, Region, RegionLike,
        Type, TypeLike, Value, ValueLike,
        attribute::{
            ArrayAttribute, DenseElementsAttribute, FlatSymbolRefAttribute, FloatAttribute,
            IntegerAttribute, StringAttribute, TypeAttribute,
        },
        operation::{OperationBuilder, OperationLike},
        r#type::{FunctionType, IntegerType, RankedTensorType},
    },
};
use mlir_sys::MlirType;
use thiserror::Error;

use crate::{
    ast::{Expr, Struct},
    parser::Parser,
    rust_side::{
        ScopedHashTable, ScopedHashTableScope, mlir_declare, mlir_get_struct_element_type,
        mlir_lookup, mlir_set_private, mlir_struct_type_from,
    },
};

#[derive(Debug, Error)]
pub enum CompileError {
    #[error("{0}")]
    ParseError(#[from] crate::parser::ParseError),
    #[error("{0}")]
    MeliorError(#[from] melior::Error),
    #[error("definition of {0} not found")]
    NotFound(String),
    #[error("{0}")]
    Utf8Error(#[from] std::str::Utf8Error),
    #[error("{0}")]
    IntoStringError(#[from] std::ffi::IntoStringError),
    #[error("struct {0} already exist")]
    StructAlreadyExist(String),
    #[error("struct {0} has init value error")]
    StructDefinitionHasInitValue(String),
    #[error("unknown struct type: {0}")]
    UnknowStructType(String),
    #[error("missing initializer")]
    MissingInitializer,
    #[error("not same type")]
    NotSameType,
    #[error("invalid access")]
    InvalidAccess,
    #[error("invalid element type")]
    InvalidElementType,
}

pub type Result<T> = std::result::Result<T, CompileError>;

pub struct Compiler {
    scoped_hash_table: ScopedHashTable,
    struct_map: HashMap<String, (MlirType, Struct)>,
}

impl Compiler {
    pub fn new() -> Self {
        Self {
            scoped_hash_table: ScopedHashTable::new(),
            // function_map: HashMap::new(),
            struct_map: HashMap::new(),
        }
    }

    pub fn compile<'c, P: AsRef<Path>>(
        &mut self,
        context: &'c Context,
        path: P,
    ) -> Result<Module<'c>> {
        let mut parser = Parser::new(path)?;
        let m = parser.parse()?;
        self.compile_module(context, &m)
    }

    fn compile_module<'c>(
        &mut self,
        context: &'c Context,
        m: &crate::ast::Module,
    ) -> Result<Module<'c>> {
        let location = Location::unknown(context);

        let module = Module::new(location);
        for record in &m.records {
            match record {
                crate::ast::Record::Function(function) => {
                    self.compile_function(context, &module, function)?
                }
                crate::ast::Record::Struct(r#struct) => {
                    self.compile_struct(context, &module, r#struct)?
                }
            }
        }
        assert!(module.as_operation().verify());
        Ok(module)
    }

    fn compile_function(
        &self,
        context: &Context,
        module: &Module,
        func: &crate::ast::Function,
    ) -> Result<()> {
        let location = Location::unknown(context);

        let _scope = ScopedHashTableScope::new(&self.scoped_hash_table);

        let mut func_type: FunctionType = self.compile_prototype(context, &func.proto)?;

        let mut arguments = Vec::new();
        for index in 0..func_type.input_count() {
            arguments.push((func_type.input(index)?, location));
        }
        let block = Block::new(&arguments);
        let body = Region::new();
        let block_ref = body.append_block(block);
        for index in 0..block_ref.argument_count() {
            let arg = block_ref.argument(index).unwrap();
            let var = &func.proto.args[index];
            mlir_declare(&self.scoped_hash_table, var, arg.into());
        }

        self.compile_expr_list(context, block_ref, &func.body)?;

        let return_op_ref = match block_ref.terminator() {
            Some(terminator) => {
                (terminator.name().as_string_ref().as_str()? == "toy.return").then_some(terminator)
            }
            None => None,
        };

        let sym_name = StringAttribute::new(context, &func.proto.name);
        match return_op_ref {
            Some(return_op_ref) => {
                let mut inputs = Vec::new();
                for index in 0..func_type.input_count() {
                    inputs.push(func_type.input(index).unwrap());
                }

                let results = return_op_ref
                    .operands()
                    .map(|operand| operand.r#type())
                    .collect::<Vec<Type>>();

                func_type = FunctionType::new(context, &inputs, &results);
            }
            None => {
                block_ref.append_operation(OperationBuilder::new("toy.return", location).build()?);
            }
        }

        let function_type = TypeAttribute::new(func_type.into());
        let func_op = OperationBuilder::new("toy.func", location)
            .add_attributes(&[
                (Identifier::new(context, "sym_name"), sym_name.into()),
                (
                    Identifier::new(context, "function_type"),
                    function_type.into(),
                ),
            ])
            .add_regions([body])
            .build()?;
        let func_op_ref = module.body().append_operation(func_op.into());
        if &func.proto.name != "main" {
            mlir_set_private(func_op_ref);
        }
        Ok(())
    }

    fn compile_prototype<'c>(
        &self,
        context: &'c Context,
        prototype: &crate::ast::Prototype,
    ) -> Result<FunctionType<'c>> {
        let mut inputs = Vec::new();
        for arg in &prototype.args {
            inputs.push(self.get_type_from_var_type(context, &arg.r#type)?);
        }

        let results = &[];
        Ok(FunctionType::new(context, &inputs, results))
    }

    fn get_type<'c>(&self, context: &'c Context, dimensions: &[u64]) -> Result<Type<'c>> {
        let f64_type = Type::float64(context);
        Ok(if dimensions.is_empty() {
            unsafe { Type::from_raw(mlir_sys::mlirUnrankedTensorTypeGet(f64_type.to_raw())) }
        } else {
            RankedTensorType::new(dimensions, f64_type, None).into()
        })
    }

    fn compile_expr_list(
        &self,
        context: &Context,
        block_ref: BlockRef,
        expr_list: &crate::ast::ExprList,
    ) -> Result<()> {
        let _scope = ScopedHashTableScope::new(&self.scoped_hash_table);

        for expr in &expr_list.exprs {
            self.compile_expr(context, block_ref, expr)?;
        }
        Ok(())
    }

    fn compile_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        expr: &crate::ast::Expr,
    ) -> Result<Option<Value<'c, 'c>>> {
        Ok(match expr {
            crate::ast::Expr::NumberExpr(number_expr) => {
                let value = self.compile_number_expr(context, block_ref, number_expr)?;
                Some(value)
            }
            crate::ast::Expr::LiteralExpr(literal_expr) => {
                let value = self.compile_literal_expr(context, block_ref, literal_expr)?;
                Some(value)
            }
            crate::ast::Expr::VariableExpr(variable_expr) => {
                let value = self.compile_variable_expr(context, block_ref, variable_expr)?;
                Some(value)
            }
            crate::ast::Expr::VarDeclExpr(var_decl_expr) => {
                self.compile_var_decl_expr(context, block_ref, var_decl_expr)?;
                None
            }
            crate::ast::Expr::ReturnExpr(return_expr) => {
                self.compile_return_expr(context, block_ref, return_expr)?;
                None
            }
            crate::ast::Expr::BinaryExpr(binary_expr) => {
                let value = self.compile_binary_expr(context, block_ref, binary_expr)?;
                Some(value)
            }
            crate::ast::Expr::CallExpr(call_expr) => {
                let value = self.compile_call_expr(context, block_ref, call_expr)?;
                Some(value)
            }
            crate::ast::Expr::PrintExpr(print_expr) => {
                self.compile_print_expr(context, block_ref, print_expr)?;
                None
            }
            crate::ast::Expr::StructLiteralExpr(struct_literal_expr) => {
                let value =
                    self.compile_struct_literal_expr(context, block_ref, struct_literal_expr)?;
                Some(value)
            }
        })
    }

    fn compile_number_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        number_expr: &crate::ast::NumberExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let f64_type = Type::float64(context);
        let r#type = RankedTensorType::new(&[], f64_type, None);
        let values = vec![FloatAttribute::new(context, f64_type, number_expr.val).into()];
        let value = DenseElementsAttribute::new(r#type.into(), &values)?;
        let constant_op = OperationBuilder::new("toy.constant", location)
            .add_results(&[r#type.into()])
            .add_attributes(&[(Identifier::new(context, "value"), value.into())])
            .build()?;
        let constant_op_ref = block_ref.append_operation(constant_op);
        Ok(constant_op_ref.result(0).unwrap().into())
    }

    fn compile_literal_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        literal_expr: &crate::ast::LiteralExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let ty = self.get_type(context, &literal_expr.dims)?;
        let value = self.get_constant_attr_for_literal_expr(context, literal_expr)?;

        let constant_op = OperationBuilder::new("toy.constant", location)
            .add_results(&[ty])
            .add_attributes(&[(Identifier::new(context, "value"), value.into())])
            .build()?;
        let constant_op_ref = block_ref.append_operation(constant_op);
        let result = constant_op_ref.result(0).unwrap().into();
        Ok(result)
    }

    fn compile_variable_expr<'c>(
        &self,
        _context: &'c Context,
        _block_ref: BlockRef<'c, 'c>,
        variable_expr: &crate::ast::VariableExpr,
    ) -> Result<Value<'c, 'c>> {
        match mlir_lookup(&self.scoped_hash_table, &variable_expr.name) {
            Some((value, _var)) => Ok(value),
            None => Err(CompileError::NotFound(variable_expr.name.clone())),
        }
    }

    fn compile_var_decl_expr(
        &self,
        context: &Context,
        block_ref: BlockRef,
        var_decl_expr: &crate::ast::VarDeclExpr,
    ) -> Result<()> {
        let location = Location::unknown(context);

        if var_decl_expr.init_val.is_none() {
            return Err(CompileError::MissingInitializer);
        }

        let mut value = self
            .compile_expr(context, block_ref, var_decl_expr.init_val.as_ref().unwrap())?
            .unwrap();

        if !var_decl_expr.r#type.name.is_empty() {
            let r#type = self.get_type_from_var_type(context, &var_decl_expr.r#type)?;
            if r#type != value.r#type() {
                return Err(CompileError::NotSameType);
            }
        } else if !var_decl_expr.r#type.shape.is_empty() {
            let results = &[self.get_type(context, &var_decl_expr.r#type.shape)?];
            let operands = &[value];
            let reshape_op = OperationBuilder::new("toy.reshape", location)
                .add_operands(operands)
                .add_results(results)
                .build()?;
            let reshape_op_ref = block_ref.append_operation(reshape_op);
            value = reshape_op_ref.result(0).unwrap().into();
        }

        assert!(mlir_declare(&self.scoped_hash_table, &var_decl_expr, value));
        Ok(())
    }

    fn compile_return_expr(
        &self,
        context: &Context,
        block_ref: BlockRef,
        return_expr: &crate::ast::ReturnExpr,
    ) -> Result<()> {
        let location = Location::unknown(context);

        let input = match &return_expr.expr {
            Some(expr) => vec![self.compile_expr(context, block_ref, &*expr)?.unwrap()],
            None => Vec::new(),
        };
        let return_op = OperationBuilder::new("toy.return", location)
            .add_operands(&input)
            .build()?;
        block_ref.append_operation(return_op.into());
        Ok(())
    }

    fn compile_binary_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        binary_expr: &crate::ast::BinaryExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let lhs = Box::new(self.compile_expr(context, block_ref, &*binary_expr.lhs)?).unwrap();
        let op_ref = match binary_expr.op {
            crate::ast::BinaryOperator::Add => {
                let rhs =
                    Box::new(self.compile_expr(context, block_ref, &*binary_expr.rhs)?).unwrap();
                let results = &[self.get_type(context, &[])?];
                let add_op = OperationBuilder::new("toy.add", location)
                    .add_results(results)
                    .add_operands(&[lhs, rhs])
                    .build()?;
                block_ref.append_operation(add_op)
            }
            crate::ast::BinaryOperator::Mul => {
                let rhs =
                    Box::new(self.compile_expr(context, block_ref, &*binary_expr.rhs)?).unwrap();
                let results = &[self.get_type(context, &[])?];
                let mul_op = OperationBuilder::new("toy.mul", location)
                    .add_results(results)
                    .add_operands(&[lhs, rhs])
                    .build()?;
                block_ref.append_operation(mul_op)
            }
            crate::ast::BinaryOperator::Dot => {
                let access_index = self
                    .get_member_index(binary_expr)
                    .ok_or(CompileError::InvalidAccess)?;

                let ty = lhs.r#type();
                let element_type = mlir_get_struct_element_type(ty, access_index)
                    .ok_or(CompileError::InvalidElementType)?;
                let results = &[element_type];
                let operands = &[lhs];
                let struct_access_op = OperationBuilder::new("toy.struct_access", location)
                    .add_results(results)
                    .add_operands(operands)
                    .add_attributes(&[(
                        Identifier::new(context, "index"),
                        IntegerAttribute::new(
                            IntegerType::new(context, 64).into(),
                            access_index as i64,
                        )
                        .into(),
                    )])
                    .build()?;
                block_ref.append_operation(struct_access_op)
            }
        };
        Ok(op_ref.result(0).unwrap().into())
    }

    fn compile_call_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        call_expr: &crate::ast::CallExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let callee = call_expr.callee.clone();
        let op_ref = if callee == "transpose" {
            assert!(call_expr.args.len() == 1);

            let input = self
                .compile_expr(context, block_ref, &call_expr.args[0])?
                .unwrap();
            let results = &[self.get_type(context, &[])?];
            let operands = &[input];
            let op = OperationBuilder::new("toy.transpose", location)
                .add_results(results)
                .add_operands(operands)
                .build()?;

            block_ref.append_operation(op)
        } else {
            let function = FlatSymbolRefAttribute::new(context, &callee);
            let mut arguments = Vec::new();
            for arg in &call_expr.args {
                arguments.push(self.compile_expr(context, block_ref, arg)?.unwrap());
            }
            let result_types = &[self.get_type(context, &[])?];
            let op = OperationBuilder::new("toy.generic_call", location)
                .add_attributes(&[(Identifier::new(context, "callee"), function.into())])
                .add_operands(&arguments)
                .add_results(result_types)
                .build()?;
            block_ref.append_operation(op.into())
        };
        let result = op_ref.result(0).unwrap().into();
        Ok(result)
    }

    fn compile_print_expr(
        &self,
        context: &Context,
        block_ref: BlockRef,
        print_expr: &crate::ast::PrintExpr,
    ) -> Result<()> {
        let location = Location::unknown(context);

        let input = self
            .compile_expr(context, block_ref, &print_expr.arg)?
            .unwrap();
        let print_op = OperationBuilder::new("toy.print", location)
            .add_operands(&[input])
            .build()?;
        block_ref.append_operation(print_op.into());
        Ok(())
    }

    fn collect_expr_values(&self, expr: &crate::ast::Expr, values: &mut Vec<f64>) -> Result<()> {
        match expr {
            crate::ast::Expr::NumberExpr(number_expr) => values.push(number_expr.val),
            crate::ast::Expr::LiteralExpr(literal_expr) => {
                self.collect_literal_expr_values(literal_expr, values)?
            }
            _ => todo!(),
        }
        Ok(())
    }

    fn collect_literal_expr_values(
        &self,
        literal_expr: &crate::ast::LiteralExpr,
        values: &mut Vec<f64>,
    ) -> Result<()> {
        for expr in &literal_expr.values {
            self.collect_expr_values(expr, values)?
        }
        Ok(())
    }

    fn compile_struct(
        &mut self,
        context: &Context,
        _module: &Module,
        r#struct: &Struct,
    ) -> Result<()> {
        if self.struct_map.contains_key(&r#struct.name) {
            return Err(CompileError::StructAlreadyExist(r#struct.name.clone()));
        }

        let mut element_types = Vec::new();
        for variable in &r#struct.variables {
            if variable.init_val.is_some() {
                return Err(CompileError::StructDefinitionHasInitValue(
                    r#struct.name.clone(),
                ));
            }

            if !variable.r#type.shape.is_empty() {
                return Err(CompileError::StructDefinitionHasInitValue(
                    r#struct.name.clone(),
                ));
            }

            let r#type: melior::ir::Type =
                self.get_type_from_var_type(context, &variable.r#type)?;
            element_types.push(r#type);
        }

        let r#type = mlir_struct_type_from(element_types);
        self.struct_map
            .insert(r#struct.name.clone(), (r#type.to_raw(), r#struct.clone()));
        Ok(())
    }

    fn get_type_from_var_type<'c>(
        &self,
        context: &'c Context,
        r#type: &crate::ast::VarType,
    ) -> Result<Type<'c>> {
        if !r#type.name.is_empty() {
            let pair = self
                .struct_map
                .get(&r#type.name)
                .ok_or(CompileError::UnknowStructType(r#type.name.clone()))?;
            let ty = pair.0;
            return Ok(unsafe { Type::from_raw(ty) });
        }

        self.get_type(context, &r#type.shape)
    }

    fn compile_struct_literal_expr<'c>(
        &self,
        context: &'c Context,
        block_ref: BlockRef<'c, 'c>,
        struct_literal_expr: &crate::ast::StructLiteralExpr,
    ) -> Result<Value<'c, 'c>> {
        let location = Location::unknown(context);

        let (data_attr, data_type) =
            self.get_constant_attr_for_struct_literal_expr(context, struct_literal_expr)?;
        let struct_constant_op = OperationBuilder::new("toy.struct_constant", location)
            .add_results(&[data_type])
            .add_attributes(&[(Identifier::new(context, "value"), data_attr)])
            .build()?;
        let struct_constant_op_ref = block_ref.append_operation(struct_constant_op);
        let result = struct_constant_op_ref.result(0).unwrap().into();
        Ok(result)
    }

    fn get_constant_attr_for_struct_literal_expr<'c>(
        &self,
        context: &'c Context,
        struct_literal_expr: &crate::ast::StructLiteralExpr,
    ) -> Result<(Attribute<'c>, Type<'c>)> {
        let mut data_attr = Vec::new();
        let mut data_type = Vec::new();

        for value in &struct_literal_expr.values {
            match value {
                crate::ast::Expr::NumberExpr(number_expr) => {
                    data_attr.push(
                        self.get_constant_attr_for_number_expr(context, number_expr)?
                            .into(),
                    );
                    data_type.push(self.get_type(context, &[])?);
                }
                crate::ast::Expr::LiteralExpr(literal_expr) => {
                    data_attr.push(
                        self.get_constant_attr_for_literal_expr(context, literal_expr)?
                            .into(),
                    );
                    data_type.push(self.get_type(context, &[])?);
                }
                crate::ast::Expr::VariableExpr(_variable_expr) => todo!(),
                crate::ast::Expr::VarDeclExpr(_var_decl_expr) => todo!(),
                crate::ast::Expr::ReturnExpr(_return_expr) => todo!(),
                crate::ast::Expr::BinaryExpr(_binary_expr) => todo!(),
                crate::ast::Expr::CallExpr(_call_expr) => todo!(),
                crate::ast::Expr::PrintExpr(_print_expr) => todo!(),
                crate::ast::Expr::StructLiteralExpr(struct_literal_expr) => {
                    let (attr, ty) = self
                        .get_constant_attr_for_struct_literal_expr(context, struct_literal_expr)?;
                    data_attr.push(attr);
                    data_type.push(ty);
                }
            }
        }
        let data_attr = ArrayAttribute::new(context, &data_attr).into();
        let data_type = mlir_struct_type_from(data_type);
        Ok((data_attr, data_type))
    }

    fn get_constant_attr_for_number_expr<'c>(
        &self,
        context: &'c Context,
        number_expr: &crate::ast::NumberExpr,
    ) -> Result<DenseElementsAttribute<'c>> {
        let f64_type = Type::float64(context);
        let r#type = RankedTensorType::new(&[], f64_type, None).into();
        let values = vec![FloatAttribute::new(context, f64_type, number_expr.val).into()];
        Ok(DenseElementsAttribute::new(r#type, &values)?)
    }

    fn get_constant_attr_for_literal_expr<'c>(
        &self,
        context: &'c Context,
        literal_expr: &crate::ast::LiteralExpr,
    ) -> Result<DenseElementsAttribute<'c>> {
        let f64_type = Type::float64(context);
        let r#type = RankedTensorType::new(&literal_expr.dims, f64_type, None).into();
        let mut values = Vec::new();
        self.collect_literal_expr_values(literal_expr, &mut values)?;
        let values = values
            .iter()
            .map(|value| FloatAttribute::new(context, f64_type, *value).into())
            .collect::<Vec<Attribute>>();
        Ok(DenseElementsAttribute::new(r#type, &values)?)
    }

    fn get_member_index(&self, binary_expr: &crate::ast::BinaryExpr) -> Option<usize> {
        if let Some(r#struct) = self.get_struct_for(&binary_expr.lhs) {
            match &*binary_expr.rhs {
                crate::ast::Expr::VariableExpr(variable_expr) => r#struct
                    .variables
                    .iter()
                    .position(|variable| variable.name == variable_expr.name),
                _ => todo!(),
            }
        } else {
            None
        }
    }

    fn get_struct_for(&self, expr: &Expr) -> Option<Struct> {
        match expr {
            Expr::VariableExpr(variable_expr) => {
                mlir_lookup(&self.scoped_hash_table, &variable_expr.name)
                    .and_then(|(_, var)| Some(var.r#type.name))
            }
            Expr::BinaryExpr(binary_expr) => match &*binary_expr.rhs {
                Expr::VariableExpr(variable_expr) => self
                    .get_struct_for(&binary_expr.lhs)
                    .and_then(|parent_struct| {
                        parent_struct
                            .variables
                            .iter()
                            .find(|variable| variable.name == variable_expr.name)
                            .and_then(|decl| Some(decl.r#type.name.clone()))
                    }),
                _ => None,
            },
            _ => None,
        }
        .and_then(|struct_name| {
            if struct_name.is_empty() {
                None
            } else {
                Some(struct_name)
            }
        })
        .and_then(|struct_name| {
            self.struct_map
                .get(&struct_name)
                .and_then(|(_, value)| Some(value.clone()))
        })
    }
}
